# -*- coding: utf-8 -*-
import re
import math
import multiprocessing
import datetime

class read_prmtop():
    def __init__(self):
#        self.prmtop_Name = raw_input('please input the prmtop name = ')
        self.prmtop_Name = '1emv1.prmtop'
        prmtop = open(self.prmtop_Name,'r')
        self.charge_Matrix = []
        self.atom_Name = []
        self.residue = []
        self.IAC = []
        self.ICO = []
        self.CN1 = []
        self.CN2 = []
        self.NATOMS = 0
        self.NTYPES = 0
        while True :
            m = prmtop.readline()
            if m[0:14] == '%FLAG POINTERS' :
                prmtop.readline()
                m = prmtop.readline()
                self.NATOMS = m[:8] ; self.NATOMS = int(self.NATOMS)
                self.NTYPES = m[8:16] ; self.NTYPES = int(self.NTYPES)
                m = prmtop.readline()
                self.NRES = m[8:16] ; self.NRES = int(self.NRES)
            elif m[0:15] == '%FLAG ATOM_NAME' :
                prmtop.readline()
                nLines = self.NATOMS/20
                if nLines*20 < self.NATOMS :
                    nLines += 1
                for i in range(nLines):
                    m = prmtop.readline()
                    for j in range(20):
                        self.atom_Name.append(m[4*j:4*j+4])
            elif m[0:12] == '%FLAG CHARGE' :
                prmtop.readline()
                nLines = self.NATOMS/5
                if nLines*5 < self.NATOMS :
                    nLines += 1
                for i in range(nLines):
                    m = prmtop.readline()
                    for j in range(5):
                        self.charge_Matrix.append(m[16*j:16*j+16])   
            elif m[0:21] == '%FLAG ATOM_TYPE_INDEX' :
                prmtop.readline()
                nLines = self.NATOMS/10
                if nLines*10 < self.NATOMS :
                    nLines += 1
                for i in range(nLines):
                    m = prmtop.readline()
                    for j in range(10):
                        self.IAC.append(m[8*j:8*j+8])
            elif m[0:26] == '%FLAG NONBONDED_PARM_INDEX' :
                prmtop.readline()
                nLines = self.NTYPES*self.NTYPES/10
                if nLines*10 < self.NTYPES*self.NTYPES :
                    nLines += 1
                for i in range(nLines):
                    m = prmtop.readline()
                    for j in range(10):
                        self.ICO.append(m[8*j:8*j+8])
            elif m[0:21] == '%FLAG RESIDUE_POINTER' :
                prmtop.readline()
                nLines = self.NRES/10
                if nLines*10 < self.NTYPES :
                    nLines += 1
                for i in range(nLines):
                    m = prmtop.readline()
                    for j in range(10):
                        self.residue.append(m[8*j:8*j+8])
            elif m[0:25] == '%FLAG LENNARD_JONES_ACOEF' :
                prmtop.readline()
                nLines = self.NTYPES*(self.NTYPES+1)/2/5
                if nLines*5 < self.NTYPES*(self.NTYPES+1)/2 :
                    nLines += 1
                for i in range(nLines):
                    m = prmtop.readline()
                    for j in range(5):
                        self.CN1.append(m[16*j:16*j+16])
            elif m[0:25] == '%FLAG LENNARD_JONES_BCOEF' :
                prmtop.readline()
                nLines = self.NTYPES*(self.NTYPES+1)/2/5
                if nLines*5 < self.NTYPES*(self.NTYPES+1)/2 :
                    nLines += 1
                for i in range(nLines):
                    m = prmtop.readline()
                    for j in range(5):
                        self.CN2.append(m[16*j:16*j+16])
                if len(self.charge_Matrix) > 0 :
                    if len(self.atom_Name) > 0 :
                        if len(self.IAC) > 0 :
                            if len(self.ICO) > 0:
                                if len(self.CN1) > 0 :
                                    break
                                
    def charge(self,atomN):
        charge_atomN = self.charge_Matrix[atomN-1]
        charge_atomN = float(charge_atomN)
        return charge_atomN

    def vanAB(self,atomI,atomJ):        
        IACI = int(self.IAC[atomI-1])
        IACJ = int(self.IAC[atomJ-1])
        ICO_Index = self.ICO[self.NTYPES*(IACI-1)+IACJ-1]
        ICO_Index = int(ICO_Index)
        vanA = self.CN1[ICO_Index-1] ; vanA = float(vanA)
        vanB = self.CN2[ICO_Index-1] ; vanB = float(vanB)
        vanAB = [vanA,vanB]
        return vanAB
    
    def name(self,atomN):
        name = self.atom_Name[atomN-1]
        name = name.replace(' ','')
        return name

    def residues(self,resid_nu):
        begin = int(self.residue[resid_nu-1])
        end   = int(self.residue[resid_nu]) - 1
        return (begin,end)


